#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <mp4v2/mp4v2.h>
#include <mp4v2/general.h>

#include <fstream>

class H264toMp4
{
public:
    H264toMp4(){}
    // Create a MP4 file, the file handle save in this->pMp4Handle.
    int CreateMp4File(const char *filename);
    // Set some properties of the MP4 file.
    int SetMp4File(int w, int h, int fps, int timescale);
    // Input a packet to pMp4Handle
    int InputPacket(unsigned char* packet, int pktlen, int w, int h);

    // Input a set of nal unit, and get some nal unit.
    int GetNalu(unsigned char* packetSet, unsigned char* pNalu);
    
    int SwitchH264ToMp4(FILE* pFile, const char* filename,  int w, int h, int fps, int timeScale);

    // Close Mp4 file
    int CloseFile();

private:
    MP4FileHandle pMp4Handle;
    int width;
    int height;
    int frameRate;
    int timeScale;
    int addStream;

private:
    int GetNalu(FILE* pFile, unsigned char *pNalu);
    int InputPcaket(FILE* pFile, int w, int h);
};

int H264toMp4::CreateMp4File(const char *filename)
{
    std::ifstream f(filename);
    //if(!f.good())
    //{
        pMp4Handle = MP4Create(filename, 0);
        if(MP4_INVALID_FILE_HANDLE == pMp4Handle)
        {
            return 0;
        }
        printf("create success\n");
        return 1;
    //}
    //printf("An Exist Mp4 file\n");
    //return 0;
}

int H264toMp4::SetMp4File(int w, int h, int fps, int timescale)
{
    if (w == 0 || h == 0 || fps == 0 || timescale == 0)
    {
        printf("Mp4 setup error\n");
        return 0;
    }
    width = w;
    height = h;
    frameRate = fps;
    timeScale = timescale;
    addStream = 1;
    printf("MP4 file:\n");
    printf("width = %d, height = %d, framerate = %d, timeScale = %d\n", width, height, frameRate, timeScale);
    MP4SetTimeScale(pMp4Handle, timeScale);

    return 1;
}

int H264toMp4::InputPacket(unsigned char* packet, int pktlen, int w, int h)
{
    if(pktlen <= 0)
    {
        printf("Packet input error\n");
        return 0;
    }

    unsigned char *pBuf = (unsigned char *)malloc(w*h);
    while(1)
    {
        int len = GetNalu(packet, pBuf);
        if(len < 0)
        {
            break;
        }
        if (pBuf[0] != 0 || pBuf[1] != 0 || pBuf[2] != 0 || pBuf[3] != 1)
        {
            printf("Nalu head type error\n");
        }
        len -= 4;
        unsigned char *pNalu = pBuf + 4;
        unsigned char naluTyle = pNalu[0] & 0x1F;
        MP4TrackId videoId;

        switch (naluTyle)
        {
        case 0x07:
            printf("------------------------------------\n");
            printf("sps(%d)\n", len);
            if (addStream)
            {
                videoId = MP4AddH264VideoTrack
                        (pMp4Handle, 
                        timeScale,              // 一秒钟多少timescale
                        timeScale/frameRate,    // 每个帧有多少个timescale
                        width,                  // width
                        height,                 // height
                        pNalu[1],               // sps[1] AVCProfileIndication
                        pNalu[2],               // sps[2] profile_compat
                        pNalu[3],               // sps[3] AVCLevelIndication
                        3);                     // 4 bytes length before each NAL unit
                if (videoId == MP4_INVALID_TRACK_ID)
                {
                    printf("Error:Can't add track.\n");
                    return -1;
                }
                
                MP4SetVideoProfileLevel(pMp4Handle, 0x7F);

                addStream = 0;
            }
            MP4AddH264SequenceParameterSet(pMp4Handle, videoId, pNalu, len);
            break;

        case 0x08:
            printf("pps(%d)\n", len);
            MP4AddH264PictureParameterSet(pMp4Handle, videoId, pNalu, len);
            break;
        
        default:
            printf("slice(%d)\n", len);
            pBuf[0] = (len>>24) & 0xFF;
            pBuf[1] = (len>>16) & 0xFF;
            pBuf[2] = (len>>8) & 0xFF;
            pBuf[3] = (len>>0) & 0xFF;

            MP4WriteSample(pMp4Handle, videoId, pBuf, len+4, MP4_INVALID_DURATION, 0, 1);

            break;
        }
    }
    free(pBuf); 
    return 1;
}

int H264toMp4::InputPcaket(FILE* pFile, int w, int h)
{
    if(!pFile)
    {
        printf("InputPcaket input file error\n");
        return 0;
    }
    unsigned char *pBuf = (unsigned char *)malloc(w*h);
    while(1)
    {
        int len = GetNalu(pFile, pBuf);
        if(len <= 0)
        {
            break;
        }
        if (pBuf[0] != 0 || pBuf[1] != 0 || pBuf[2] != 0 || pBuf[3] != 1)
            continue;
        
        len -= 4;
        unsigned char *pNalu = pBuf + 4;
        unsigned char naluTyle = pNalu[0] & 0x1F;
        MP4TrackId videoId;

        switch (naluTyle)
        {
        case 0x07:
            printf("------------------------------------\n");
            printf("sps(%d)\n", len);
            if (addStream)
            {
                videoId = MP4AddH264VideoTrack
                        (pMp4Handle, 
                        timeScale,              // 一秒钟多少timescale
                        timeScale/frameRate,    // 每个帧有多少个timescale
                        width,                  // width
                        height,                 // height
                        pNalu[1],               // sps[1] AVCProfileIndication
                        pNalu[2],               // sps[2] profile_compat
                        pNalu[3],               // sps[3] AVCLevelIndication
                        3);                     // 4 bytes length before each NAL unit
                if (videoId == MP4_INVALID_TRACK_ID)
                {
                    printf("Error:Can't add track.\n");
                    return 0;
                }
                
                MP4SetVideoProfileLevel(pMp4Handle, 0x7F);

                addStream = 0;
            }
            MP4AddH264SequenceParameterSet(pMp4Handle, videoId, pNalu, len);
            break;

        case 0x08:
            printf("pps(%d)\n", len);
            MP4AddH264PictureParameterSet(pMp4Handle, videoId, pNalu, len);
            break;
        
        default:
            printf("slice(%d)\n", len);
            pBuf[0] = (len>>24) & 0xFF;
            pBuf[1] = (len>>16) & 0xFF;
            pBuf[2] = (len>>8) & 0xFF;
            pBuf[3] = (len>>0) & 0xFF;

            MP4WriteSample(pMp4Handle, videoId, pBuf, len+4, MP4_INVALID_DURATION, 0, 1);

            break;
        }
    }
    free(pBuf);
    return 1;
}

int H264toMp4::GetNalu(unsigned char* packetSet, unsigned char* pNalu)
{
    if(packetSet == NULL)
    {
        printf("GetNalu input error\n");
        return 0;
    }
    unsigned char* p = packetSet;
    memcpy(pNalu, packetSet, 4);
    if(pNalu[0] != 0 || pNalu[1] != 0 || pNalu[2] != 0 || pNalu[3] != 1)
    {
        printf("GetNalu packetSet input type error\n");
        return 0;
    }

    int pos = 4;
    p = p + 4;
    while(p)
    {
        if(*p == 0 && *(p + 1) == 0 && (p + 2) == 0 && *(p + 3) == 1)
        {
            packetSet = p;
        }
        pNalu[pos] = p[0];
        ++p;
        ++pos;
    }

    return pos + 1; 
}

int H264toMp4::GetNalu(FILE* pFile, unsigned char *pNalu)
{

    unsigned char c;
    int pos = 0;
    int len;

    if(!pFile)
        return -1;

    if((len = fread(pNalu, 1, 4, pFile)) <= 0)
        return -1;

    if(pNalu[0] != 0 || pNalu[1] != 0 || pNalu[2] != 0 || pNalu[3] != 1)
        return -1;

    pos = 4;

    while(1)
    {
        if(feof(pFile))
            break;

        pNalu[pos] = fgetc(pFile);
        
        if(pNalu[pos-3] == 0 && pNalu[pos-2] == 0 && pNalu[pos-1] == 0 && pNalu[pos] == 1)
        {
            fseek(pFile, -4, SEEK_CUR);
            pos -= 4;
            break;
        }

        pos++;
    }

    len = pos+1;

    return len;
}

int H264toMp4::CloseFile()
{
    if(pMp4Handle == NULL)
    {
        printf("close error\n");
        return 0;
    }
    MP4Close(pMp4Handle, 0);
    return 1;
}

int H264toMp4::SwitchH264ToMp4(FILE* pIn, const char* filename, int w, int h, int fps, int timeScale)
{
    if(!pIn)
    {
        printf("input file error\n");
        return 0;
    }
    int ret = 1;
    ret = CreateMp4File(filename);
    if(!ret)
    {
        printf("Create Error");
        return ret;
    }

    ret = SetMp4File(w, h, fps, timeScale);
    if(!ret)
    {
        printf("Set error\n");
        return ret;
    }

    ret = InputPcaket(pIn, w, h);
    if(!ret)
    {
        printf("error!\n");
        return ret;
    }

    ret = CloseFile();
    if(!ret)
    {
        printf("File close error\n");
        return ret;
    }
    return ret;
}